# encoding:utf-8

import json
import os

import plugins
from bridge.bridge import Bridge
from bridge.context import ContextType
from bridge.reply import Reply, ReplyType
from common import const
from common.log import logger
from config import conf
from plugins import *


class RolePlay:
    def __init__(self, bot, sessionid, desc, wrapper=None, roleName=""):
        self.bot = bot
        self.roleName = roleName
        self.sessionid = sessionid
        self.wrapper = wrapper or "%s"  # 用于包装用户输入
        self.desc = desc
        self.bot.sessions.build_session(self.sessionid, system_prompt=self.desc)

    def reset(self):
        self.bot.sessions.clear_session(self.sessionid)

    def action(self, user_action):
        session = self.bot.sessions.build_session(self.sessionid)
        if session.system_prompt != self.desc:  # 目前没有触发session过期事件，这里先简单判断，然后重置
            session.set_system_prompt(self.desc)
        prompt = self.wrapper % user_action
        return prompt


@plugins.register(
    name="Assistant",
    desire_priority=0,
    namecn="AI助理",
    desc="可以直接按场景调用",
    version="1.0",
    author="lanvent",
)
class Assistant(Plugin):
    def __init__(self):
        super().__init__()
        curdir = os.path.dirname(__file__)
        config_path = os.path.join(curdir, "assistant.json")
        try:
            with open(config_path, "r", encoding="utf-8") as f:
                config = json.load(f)
                self.tags = {tag: (desc, []) for tag, desc in config["tags"].items()}
                self.roles = {}
                for role in config["roles"]:
                    self.roles[role["title"].lower()] = role
                    for tag in role["tags"]:
                        if tag not in self.tags:
                            logger.warning(f"[Assistant] unknown tag {tag} ")
                            self.tags[tag] = (tag, [])
                        self.tags[tag][1].append(role)
                for tag in list(self.tags.keys()):
                    if len(self.tags[tag][1]) == 0:
                        logger.debug(f"[Assistant] no role found for tag {tag} ")
                        del self.tags[tag]

            if len(self.roles) == 0:
                raise Exception("no role found")
            self.handlers[Event.ON_HANDLE_CONTEXT] = self.on_handle_context
            self.roleplays = {}
            logger.info("[Assistant] inited")
        except Exception as e:
            if isinstance(e, FileNotFoundError):
                logger.warn(f"[Assistant] init failed, {config_path} not found, ignore or see https://github.com/zhayujie/chatgpt-on-wechat/tree/master/plugins/role .")
            else:
                logger.warn("[Assistant] init failed, ignore or see https://github.com/zhayujie/chatgpt-on-wechat/tree/master/plugins/role .")
            raise e


    def get_role(self, name):
        name = name.lower()
        found_role = None
        if name in self.roles:
            found_role = name
        return found_role

    def on_handle_context(self, e_context: EventContext):
        logger.debug("[Assistant] on_handle_context. context: %s" % e_context.econtext)
        logger.debug("[Assistant] on_handle_context. context: %s" % str(e_context['context']))
        if e_context["context"].type != ContextType.TEXT:
            return
        btype = Bridge().get_bot_type("chat")
        if btype not in [const.OPEN_AI, const.CHATGPT, const.CHATGPTONAZURE, const.LINKAI]:
            return
        bot = Bridge().get_bot("chat")
        content = e_context["context"].content[:]
        # clist = e_context["context"].content.split(maxsplit=1)
        customize = False
        sessionid = e_context["context"]["session_id"]
        # trigger_prefix = conf().get("plugin_trigger_prefix", "$")
        trigger_prefix = "@"
        if content[:1] == trigger_prefix:
            assistant = content[1:]
            logger.debug("[Assistant] handle cmd: %s" % assistant)
            if assistant.lower() in ["help", "帮助"] or assistant is None or len(assistant) == 0:
                logger.debug("[Assistant] handle help: %s" % assistant)
                help_text = "助理列表：\n"
                for role in self.roles.values():
                    help_text += f"{role['title']}: {role['remark']}\n"
                help_text += "直接使用 @助理名称 触发， @退出助理 退出"
                reply = Reply(ReplyType.INFO, help_text)
                e_context["reply"] = reply
                e_context.action = EventAction.BREAK_PASS
                return
            elif assistant.lower() in ["退出", "退出助理"]:
                rolename = ""
                if sessionid in self.roleplays:
                    rolename = self.roleplays[sessionid].roleName
                    self.roleplays[sessionid].reset()
                    del self.roleplays[sessionid]

                reply = Reply(ReplyType.INFO, f"已退出${rolename}!")
                e_context["reply"] = reply
                e_context.action = EventAction.BREAK_PASS
                return
            role = self.get_role(assistant)
            if role is None:
                reply = Reply(ReplyType.ERROR, "助理不存在")
                e_context["reply"] = reply
                e_context.action = EventAction.BREAK_PASS
                return
            else:
                self.roleplays[sessionid] = RolePlay(
                    bot,
                    sessionid,
                    self.roles[role]["descn"],
                    self.roles[role].get("wrapper", "%s"),
                    self.roles[role]["title"],
                )
                reply = Reply(ReplyType.INFO, f"激活的助理为 {role}:\n" + self.roles[role]["title"])
                e_context["reply"] = reply
                e_context.action = EventAction.BREAK_PASS
                return
        elif sessionid not in self.roleplays:
            logger.debug("[Assistant] session not in assitant  user array")
            return


        text = f"🚀您的内容将在10-20秒左右完成，请耐心等待\n"
        text += f"- - - - - - - - -\n"
        text += content
        reply = Reply(ReplyType.INFO, text)
        channel = e_context["channel"]
        _send(channel, reply, e_context["context"])
        prompt = self.roleplays[sessionid].action(content)
        e_context["context"].type = ContextType.TEXT
        e_context["context"].content = prompt
        e_context.action = EventAction.BREAK
def _send(channel, reply: Reply, context, retry_cnt=0):
    try:
        channel.send(reply, context)
    except Exception as e:
        logger.error("[WX] sendMsg error: {}".format(str(e)))
        if isinstance(e, NotImplementedError):
            return
        logger.exception(e)
        if retry_cnt < 2:
            time.sleep(3 + 3 * retry_cnt)
            channel.send(reply, context, retry_cnt + 1)
